import requests

from config import zhipu_api_key

def calculate_tokens(message, api_key=zhipu_api_key, model="GLM-4-Air"):
    url = 'https://open.bigmodel.cn/api/paas/v4/tokenizer'
    headers = {
        'Authorization': api_key,
        'Content-Type': 'application/json'
    }
    data = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": message
            }
        ]
    }
    
    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        result = response.json()
        return result['usage']['prompt_tokens']
    else:
        print(f"请求失败，状态码：{response.status_code}")
        return None

# 使用示例
def main():
    message = """
1 对一个技术概念的调研综述
2 对一个方向研究现状的综述
3 对多篇文献的对比分析综述
4 对一个技术方法研究脉络的综述
你是一个分类专家专家，请你将用户的输入分类到上面的类别中，仅输出一个数字表示类别。

输入：损失函数，输出：1
输入：Text2SQL研究现状如何，面临哪些挑战，输出：2
输入：有哪些方法可以提升大模型的规划能力，各自优劣是什么，输出：3
输入：多模态大模型的技术发展路线是什么样的，输出：4
输入：损失函数，输出：
"""
    tokens = calculate_tokens(message)
    if tokens is not None:
        print(f"消息的token数量为：{tokens}")

# 运行示例
if __name__ == "__main__":
    main()

"""
记录全流程 tokens 数量

问题分类 150

"""